{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7a434f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "BRANCH=\"main\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "developmental-gibraltar",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "You can run either this notebook locally (if you have all the dependencies and a GPU) or on Google Colab.\n",
    "\n",
    "Instructions for setting up Colab are as follows:\n",
    "1. Open a new Python 3 notebook.\n",
    "2. Import this notebook from GitHub (File -> Upload Notebook -> \"GITHUB\" tab -> copy/paste GitHub URL)\n",
    "3. Connect to an instance with a GPU (Runtime -> Change runtime type -> select \"GPU\" for hardware accelerator)\n",
    "4. Run this cell to set up dependencies.\n",
    "\"\"\"\n",
    "# If you're using Google Colab and not running locally, run this cell\n",
    "\n",
    "# install NeMo\n",
    "!python -m pip install git+https://github.com/NVIDIA/NeMo.git@$BRANCH#egg=nemo_toolkit[nlp]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42daf8bf",
   "metadata": {},
   "source": [
    "# Introduction\n",
    "\n",
    "In this notebook we demonstrate how to use p-tunining and prompt tuning within NeMo-Megatron. Both methods are parameter efficient alternatives to fine-tuning pretrained language models. Our NeMo implementation makes it possible to use one pretrained GPT model on many downstream tasks without needing to tune the model’s full set of parameters. It also allows for adding new tasks to your model without overwriting or disrupting previous tasks for which the model has already been p-tuned/prompt-tuned. Because the original model parameters are frozen and never altered by either method, p-tuning/prompt-tuning also avoid cartographic forgetting issues often encountered when fine-tuning models.\n",
    "\n",
    "- Our prompt tuning implementation is based off Lester et. al’s EMNLP 2021 paper [The Power of Scale for Parameter-Efficient Prompt Tuning](https://arxiv.org/abs/2104.08691)\n",
    "\n",
    "- Our p-tuning implementation is based off Liu et al's paper [GPT Understands, Too](https://arxiv.org/abs/2103.10385).\n",
    "\n",
    "- Usage examples and API documentation can be found in [our user docs](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/nlp/prompt_learning.html). \n",
    "\n",
    "<img src=\"images/prompt_learning_forward_pass.png\" alt=\"Prompt Learning Forward Pass\"/>\n",
    "\n",
    "Our continuous learning capability for combined p-tuning and prompt tuning with GPT style models is a NeMo specific extension of the author’s original work.\n",
    "\n",
    "# The Plan\n",
    "\n",
    "We are going to show you how to:\n",
    "    \n",
    "    1. P-Tune/Prompt Tune a model on multiple tasks at the same time\n",
    "    2. Add a new task to a model that has already been P-Tuned/Prompt Tuned previously\n",
    "    \n",
    "We will first p-tune a GPT model on sentiment analysis, and intent and slot classification tasks. Then we will show how to add the squad question answering task to the same model we already p-tuned once.\n",
    "\n",
    "\n",
    "# Techincal Overview\n",
    "Instead of selecting discrete text prompts in a manual or automated fashion, prompt tuning and p-tuning utilize virtual prompt embeddings that can be optimized via gradient decent. The only difference between prompt tuning and p-tuning within NeMo-Megatron is the architecture used to tune the soft prompt tokens during training.\n",
    "\n",
    "### Terminology\n",
    "We will be using the terms `continuous`, `soft`, and `virtual` token interchangeably to refer to embeddings inserted into the model prompt that have no concrete mapping to strings or characters within the model’s vocabulary. These virtual token embeddings exist in contrast to the `discrete`, `hard`, or `real` tokens that do make up the model’s vocabulary. Virtual tokens are purely 1D vectors with dimensionality equal to that of each real token embedding, matching the `hidden_size` hyperparameter. In training and inference, continuous token embeddings are inserted among discrete token embeddings according to a template you provide in the model’s config. We will demonstrate how to do this below.\n",
    "\n",
    "When referring to p-tuning and prompt tuning together, we will be using the phrase prompt learning for simplicity.\n",
    "\n",
    "### Prompt-Tuning\n",
    "In prompt-tuning a pretrained GPT model, soft prompt embeddings are initialized as a 2D matrix of size `total_virtual_tokens X hidden_size`. Each task the model is prompt-tuned to perform has its own 2D embedding matrix associated with it. Tasks do not share any parameters during training or inference. All GPT model parameters are frozen and only the embedding parameters for each task are updated during training.\n",
    "\n",
    "In prompt tuning you can specify how the embeddings are initialized for each task. You can either\n",
    "\n",
    "1. Initialize embedding parameters according to some random distribution\n",
    "2. Initialize embedding parameters from existing vocabulary embeddings (recommended)\n",
    "\n",
    "If you choose to initialize virtual token embeddings from existing embedding weights, you can provide the string of words you want to use for initialization in the model’s config. This string will be tokenized and tiled or truncated to match the specified number of virtual tokens you would like to use (`total_virtual_tokens`). Vocab embeddings are copied and used to initialize the soft prompt embedding matrix for each task. The vocab embeddings themselves are not updated or changed during prompt tuning.\n",
    "\n",
    "\n",
    "### P-Tuning\n",
    "In p-tuning, an LSTM model is used to predict virtual token embeddings. We refer to this LSTM model as our `prompt_encoder`. LSTM parameters are randomly initialized at the start of p-tuning. All GPT model parameters are frozen, and only the LSTM weights are updated at each training step. LSTM parameters are shared between all tasks that are p-tuned at the same time, but the LSTM model outputs unique virtual token embeddings for each task. The virtual tokens predicted by the LSTM are inserted among the discrete token input in the exact same manner as with prompt-tuning. You still specify the number of virtual tokens you want to use by setting `total_virtual_tokens` and each virtual token embedding is still a 1D vector of size `hidden_size`.\n",
    "\n",
    "\n",
    "\n",
    "# The Best of Both\n",
    "A single pretrained GPT model can use both p-tuning and prompt-tuning. While you must decide to use either p-tuning or prompt-tuning for each task you want your model to perform, you can p-tune your model on a set of tasks A, then prompt tune your same model on a different set of tasks B, then finally run inference on tasks from both A and B at the same time. During prompt-tuning or p-tuning, tasks tuned at the same time must use the same number of virtual tokens. During inference, tasks using differing amounts of virtual tokens can be run at the same time.\n",
    "\n",
    "Please see our [docs for more comparisons between prompt and p-tuning](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/nlp/prompt_learning.html). \n",
    "\n",
    "With all that covered, let's get started!\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31c27562",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import wget \n",
    "import pathlib"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0bfc7709",
   "metadata": {},
   "source": [
    "# Tasks and Datasets\n",
    "We will be using p-tuning to teach our GPT model to do 3 tasks: **Sentiment Analysis**, **Question Answering** and **Intent and Slot Classification**.\n",
    "\n",
    "We will use [Financial PhraseBank dataset](https://huggingface.co/datasets/financial_phrasebank) for our sentiment analysis task, [SQuAD](https://rajpurkar.github.io/SQuAD-explorer/) for question answering task, and the [Assistant Benchmarking Dataset](https://github.com/xliuhw/NLU-Evaluation-Data) for intent and slot classification. \n",
    "\n",
    "- The [Financial PhraseBank dataset](https://huggingface.co/datasets/financial_phrasebank) contains the sentiments for financial news headlines from the perspective of a retail investor. Further details about the dataset can be found in Malo et. al's \"[Good Debt or Bad Debt: Detecting Semantic Orientations in Economic Texts](https://arxiv.org/abs/1307.5336)\"\n",
    "\n",
    "\n",
    "- [SQuAD](https://rajpurkar.github.io/SQuAD-explorer/) is a reading comprehension dataset, consisting of questions posed by crowd workers on a set of Wikipedia articles, where the answer to every question is a segment of text. More information on [SQuAD](https://rajpurkar.github.io/SQuAD-explorer/) can be found on their website or in their paper by Rajpurkar et. al \"[Know What You Don’t Know: Unanswerable Questions for SQuAD](https://arxiv.org/pdf/1806.03822.pdf)\".\n",
    "\n",
    "\n",
    "- The [Assistant Benchmarking Dataset](https://github.com/xliuhw/NLU-Evaluation-Data) is a natural language dataset for in home human-robot interaction. Details on the dataset can be found in Liu et. al's \"[Benchmarking Natural Language Understanding Services for building Conversational Agents](https://arxiv.org/abs/1903.05566)\"\n",
    "\n",
    "Each of these tasks require different types of natural language understanding and lie in different domains. We will demonstrate how to use one model to perform all of them."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0b0072a",
   "metadata": {},
   "source": [
    "# Data Preparation\n",
    "\n",
    "The prompt learning dataset loader accepts a list of json/dictionary objects or a list of json file names where each json file contains a collection of json objects. Each json object must include the field `taskname` which is a string identifier for the task the data example corresponds to. They should also include one or more fields corresponding to different sections of the discrete text prompt. The input data might look like:\n",
    "\n",
    "```\n",
    "[\n",
    "    {\"taskname\": \"squad\", \"context\": [CONTEXT_PARAGRAPH_TEXT1], \"question\": [QUESTION_TEXT1], \"answer\": [ANSWER_TEXT1]},\n",
    "    {\"taskname\": \"squad\", \"context\": [CONTEXT_PARAGRAPH_TEXT2], \"question\": [QUESTION_TEXT2], \"answer\": [ANSWER_TEXT2]},\n",
    "    {\"taskname\": \"intent_and_slot\", \"utterance\": [UTTERANCE_TEXT1], \"label\": [INTENT_TEXT1][SLOT_TEXT1]},\n",
    "    {\"taskname\": \"intent_and_slot\", \"utterance\": [UTTERANCE_TEXT2], \"label\": [INTENT_TEXT2][SLOT_TEXT2]},\n",
    "    {\"taskname\": \"sentiment\", \"sentence\": [SENTENCE_TEXT1], \"label\": [SENTIMENT_LABEL1]},\n",
    "    {\"taskname\": \"sentiment\", \"sentence\": [SENTENCE_TEXT2], \"label\": [SENTIMENT_LABEL2]},\n",
    "]\n",
    "```\n",
    "\n",
    "These additional fields can be unlimited in number and will be used to help map different parts of the discrete text input to a prompt template that you define. We will show how this mapping works and how to construct your prompt template in the `Prompt Formatting` section. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0dbd41fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# You can replace DATA_DIR and NEMO_DIR with your own locations\n",
    "DATA_DIR = \"data\"\n",
    "NEMO_DIR = '.'\n",
    "\n",
    "os.makedirs(DATA_DIR, exist_ok=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "504a7b40",
   "metadata": {},
   "source": [
    "\n",
    "For each dataset we have preprocessing scripts pre-written in NeMo's example directory located in `examples/nlp`. Let's download those now. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e72a1dc1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# download the preprocessing scripts from github for the purpose of this tutorial\n",
    "wget.download(f'https://raw.githubusercontent.com/NVIDIA/NeMo/{BRANCH}/scripts/dataset_processing/nlp/financial_phrase_bank/prompt_learning_financial_phrase_bank_preprocessing.py', NEMO_DIR)\n",
    "wget.download(f'https://raw.githubusercontent.com/NVIDIA/NeMo/{BRANCH}/scripts/dataset_processing/nlp/squad/prompt_learning_squad_preprocessing.py', NEMO_DIR)\n",
    "wget.download(f'https://raw.githubusercontent.com/NVIDIA/NeMo/{BRANCH}/scripts/dataset_processing/nlp/intent_and_slot/prompt_learning_assistant_preprocessing.py', NEMO_DIR)\n",
    "wget.download(f'https://raw.githubusercontent.com/NVIDIA/NeMo/{BRANCH}/scripts/dataset_processing/nlp/intent_and_slot/assistant_utils.py', NEMO_DIR)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71813919",
   "metadata": {},
   "source": [
    "Now let's down load and process each dataset."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a09c85b2",
   "metadata": {},
   "source": [
    "###  Financial PhraseBank Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7ce77b8",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Download the financial phrase bank dataset\n",
    "!wget https://www.researchgate.net/profile/Pekka_Malo/publication/251231364_FinancialPhraseBank-v10/data/0c96051eee4fb1d56e000000/FinancialPhraseBank-v10.zip\n",
    "!unzip FinancialPhraseBank-v10.zip -d {DATA_DIR}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e253e89b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# What the financial phrase bank dataset looks like before processing\n",
    "SENTIMENT_DIR = os.path.join(DATA_DIR, \"FinancialPhraseBank-v1.0\")\n",
    "!head -4 $SENTIMENT_DIR/Sentences_AllAgree.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10ec9405",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Preprocess financial phrase bank dataset\n",
    "!python $NEMO_DIR/prompt_learning_financial_phrase_bank_preprocessing.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4077d21b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# What the financial phrase bank dataset looks like after processing\n",
    "!head -4 $SENTIMENT_DIR/financial_phrase_bank_train.jsonl"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06481f49",
   "metadata": {},
   "source": [
    "Our financial phrase bank preprocessing script converted the raw text file of sentences and labels into three `.jsonl` files for training, validation, and testing. Each line in the files contains a json object with the fields `taskname`, `sentiment`,`sentence`, and `label`. You can inspect the preprocessing script and play with different arguments for the script by looking at and running `prompt_learning_financial_phrase_bank_preprocessing.py` which should currently be downloaded in `NEMO_DIR`. It is also located at `scripts/dataset_processing/nlp/financial_phrase_bank/prompt_learning_financial_phrase_bank_preprocessing.py` in the NeMo repo.\n",
    "\n",
    "By default 80% of the data was randomly selected for the training set, 10% for the validation set, and 10% for the test set. We only used training examples with 100% agreement from labelers on the correct sentiment label. This data is from `Sentences_AllAgree.txt`. This should result in `1811` training examples, `226` validation examples, and `227` examples for testing. The `label` field was removed from test examples. \n",
    "\n",
    "If you want to try using more data, you can combine the `Sentences_AllAgree.txt` with any of the `Sentences_75Agree.txt`, `Sentences_66Agree.txt` and/or `Sentences_50Agree.txt` by creating a new cell below and running:\n",
    "\n",
    "```\n",
    "!cat {SENTIMENT_DIR}/Sentences_AllAgree.txt {SENTIMENT_DIR}/Sentences_75Agree.txt >> {SENTIMENT_DIR}/combined_data.txt\n",
    "!python $NEMO_DIR/prompt_learning_financial_phrase_bank_preprocessing.py --file-name combined_data.txt\n",
    "```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "816791de",
   "metadata": {},
   "source": [
    "### SQuAD Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa16d8ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "SQUAD_DIR = os.path.join(DATA_DIR, \"SQuAD\")\n",
    "os.makedirs(SQUAD_DIR, exist_ok=True)\n",
    "\n",
    "# Download the SQuAD dataset\n",
    "!wget https://rajpurkar.github.io/SQuAD-explorer/dataset/train-v2.0.json\n",
    "!mv train-v2.0.json {SQUAD_DIR}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64e3e25b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Preprocess squad data\n",
    "!python $NEMO_DIR/prompt_learning_squad_preprocessing.py --data-dir {SQUAD_DIR}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b562d1de",
   "metadata": {},
   "outputs": [],
   "source": [
    "# What the squad dataset looks like after processing\n",
    "!head -4 $SQUAD_DIR/squad_train.jsonl"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a385d319",
   "metadata": {},
   "source": [
    "We made a `.jsonl` file for each of the train, validation, and testing splits of the squad data. Every `.jsonl` file contains json objects with the fields `taskname`, `context`, `question`, and `answer`. The preprocessing script is called `prompt_learning_squad_preprocessing.py`. It should be in your `NEMO_DIR` and at `scripts/dataset_processing/nlp/squad/prompt_learning_squad_preprocessing.py` in the NeMo repo. \n",
    "\n",
    "The SQuAD dataset consists of various topics like `Beyoncé`, `IPod`, and `Symbiosis`. Each topic has several paragraphs associated with it, and each paragraph has several questions and answers related to it. When we separated the train/validation/test splits, we separated them on the topic level. For example, if the training set contains paragraphs and questions about the topic `Beyoncé`, neither the validation nor test sets will contain any questions on this topic. All questions about a certain topic are isolated to one split of the data. \n",
    "\n",
    "Like the Financial PhraseBank Dataset, we randomly selected 80% of the questions for training, 10% for validation, and 10% for test. This resulted in `69125` test examples, `8952` validation examples, and `8744` testing examples. The `answer` field was removed from test examples.\n",
    "\n",
    "Training on the full train split could take a lot of time, so we are going to clip the train split to 20k examples for the sake of this tutorial. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f1473ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "! head -20000 $SQUAD_DIR/squad_train.jsonl > $SQUAD_DIR/squad_short_train.jsonl"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91626be3",
   "metadata": {},
   "source": [
    "### Assistant Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9340eb9",
   "metadata": {},
   "outputs": [],
   "source": [
    "ASSISTANT_DIR = os.path.join(DATA_DIR, \"assistant\")\n",
    "os.makedirs(ASSISTANT_DIR, exist_ok=True)\n",
    "\n",
    "# Download the assisent dataset\n",
    "!wget https://github.com/xliuhw/NLU-Evaluation-Data/archive/master.zip\n",
    "!unzip master.zip -d {ASSISTANT_DIR}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8194a460",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Process virtual assistant intent and slot classification data\n",
    "!python $NEMO_DIR/prompt_learning_assistant_preprocessing.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67dc4c1f",
   "metadata": {},
   "outputs": [],
   "source": [
    "!head -5 $ASSISTANT_DIR/assistant_train.jsonl"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e803eaed",
   "metadata": {},
   "source": [
    "For the virtual assistent dataset, there are a set of 64 possible intents:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "812707c8",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!echo 'Intents: ' $(wc -l < {ASSISTANT_DIR}/nemo-format/dict.intents.csv)\n",
    "!cat {ASSISTANT_DIR}/nemo-format/dict.intents.csv"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a6b7020",
   "metadata": {},
   "source": [
    "and 55 types of slots:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55507018",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# print all slots from the NeMo format slot dictionary\n",
    "!echo 'Slots: ' $(wc -l < {ASSISTANT_DIR}/nemo-format/dict.slots.csv)\n",
    "!cat {ASSISTANT_DIR}/nemo-format/dict.slots.csv"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37cd65ff",
   "metadata": {},
   "source": [
    "Each slot label consists of the slot type followed by specific text from the utterance corresponding to that slot type in parentheses. For example, the utterance `\"tell my facebook group that i've arrived\"` has the intent label `social_post` and the slot label `media_type(facebook)`. Utterances each have one intent label and zero or more slot labels. In cases where there is no slot label, our GPT model should predict the word `None`. \n",
    "\n",
    "Json objects for each training example contain three fields: `taskname`, `utterance`, and `label`. For this dataset, our preprocessing scipt formatted our intent and slot labels to look like `\"\\nIntent: transport_taxi\\nSlots: transport_agency(golden taxi), time(seven pm), date(today)\"`. With newline characters (\\n) separating intent and slot labels. Our train jsonl file has `9960` training examples. Our validation and test jsonl files each have `538` training examples. Test examples do not have the `label` field. \n",
    "\n",
    "The preprocessing script can be found at `scripts/dataset_processing/nlp/intent_and_slot/prompt_learning_assistant_preprocessing.py`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e19c8dc",
   "metadata": {},
   "source": [
    "# P-Tuning Model Config Setup\n",
    "\n",
    "Now we will begin setting up the conifg file used for prompt/p-tuning our GPT models! GPT Prompt learning within NeMo uses a class called `MegatronGPTPromptLearningModel` which has its own config file. We will start by loading an example prompt learning config file, then make changes to it to fit our tasks and training plans. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5749c387",
   "metadata": {},
   "outputs": [],
   "source": [
    "from omegaconf import OmegaConf\n",
    "\n",
    "CONFIG_DIR = os.path.join(NEMO_DIR, \"conf\")\n",
    "os.makedirs(CONFIG_DIR, exist_ok=True)\n",
    "\n",
    "# Download the example config file\n",
    "wget.download(f'https://raw.githubusercontent.com/NVIDIA/NeMo/{BRANCH}/examples/nlp/language_modeling/conf/megatron_gpt_prompt_learning_config.yaml', CONFIG_DIR)\n",
    "\n",
    "# Load the example config file so we can start editing it\n",
    "CONFIG_PATH = os.path.join(CONFIG_DIR, \"megatron_gpt_prompt_learning_config.yaml\")\n",
    "config = OmegaConf.load(CONFIG_PATH)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce966bcf",
   "metadata": {},
   "source": [
    "First let's set the datasets we've created in the config. We are going to start by p-tuning a GPT model on the **Sentiment Analysis** and **Intent and Slot Classification** tasks. We will set only those two datasets in the config file right now. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6bb1590f",
   "metadata": {},
   "outputs": [],
   "source": [
    "config.model.data.train_ds = [f\"{SENTIMENT_DIR}/financial_phrase_bank_train.jsonl\", f\"{ASSISTANT_DIR}/assistant_train.jsonl\"]\n",
    "config.model.data.validation_ds = [f\"{SENTIMENT_DIR}/financial_phrase_bank_val.jsonl\", f\"{ASSISTANT_DIR}/assistant_val.jsonl\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed04f167",
   "metadata": {},
   "source": [
    "The `MegatronGPTPromptLearningModel` class expects datasets to be a list of `.json` or `.jsonl` file paths. You can give it multiple datasets at once like we did above.  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e021b24",
   "metadata": {},
   "source": [
    "### Prompt Formatting\n",
    "Now that we have our datasets set, lets define what we want the prompt for each task to look like. \n",
    "\n",
    "To customize different prompts for different tasks, we simply need to specify the prompt task template in the config file. The virtual token markers `<|VIRTUAL_PROMPT_#|>` signify where you want virtual tokens to be placed in the template string. `<|VIRTUAL_PROMPT_0|>`, `<|VIRTUAL_PROMPT_1|>`, and `<|VIRTUAL_PROMPT_2|>` indicate where a number of virtual tokens matching the values given at `virtual_token_splits[0]`, `virtual_token_splits[1]` and `virtual_token_splits[2]` will be placed. The other variable fields `{var}` refer to the fields in the data json. \n",
    "\n",
    "For example, given: \n",
    "\n",
    "1. the data json **{\"sentence1\": \"And he said, Mama, I'm home.\", \"sentence2\": \"He didn't say a word.\"}**\n",
    "\n",
    "\n",
    "2. virtual token splits set to `virtual_token_splits = [3, 3, 3]` \n",
    "\n",
    "\n",
    "3. a prompt template set to `prompt_template = \"<|VIRTUAL_PROMPT_0|> Hypothesis: [sentence1], <|VIRTUAL_PROMPT_1|> Premise: [sentence2] <|VIRTUAL_PROMPT_2|> Answer:\"`\n",
    "\n",
    "the input will be translated into **<span style=\"color:red\">VVV</span> Hypothesis: And he said, Mama, I'm home.<span style=\"color:red\">VVV</span> Premise: He didn't say a word.<span style=\"color:red\">VVV</span> Answer:**, where <span style=\"color:red\">VVV</span> are three virtual tokens.\n",
    "\n",
    "Because we are only p-tuning on the Sentiment Analysis and Intent/Slot Classification tasks right now, we only need to set the task templates for those two tasks. But, we are going to go ahead and set the template for all 3 tasks here, just to show that you can set the templates for all the tasks you have planned at one time, then prompt tune/p-tune on them sequentially. \n",
    "\n",
    "Let's configure all of our templates for **Sentiment Analysis**, **Intent and Slot Classification**, and **Question Answering** tasks:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f935b411",
   "metadata": {},
   "outputs": [],
   "source": [
    "config.model.task_templates = [\n",
    "    {\n",
    "        \"taskname\": \"sentiment\",\n",
    "        \"prompt_template\": \"<|VIRTUAL_PROMPT_0|> {sentence} sentiment:{label}\",\n",
    "        \"total_virtual_tokens\": 10,\n",
    "        \"virtual_token_splits\": [10],\n",
    "        \"truncate_field\": None,\n",
    "        \"answer_only_loss\": True,\n",
    "        \"answer_field\": \"label\",\n",
    "    },\n",
    "    {\n",
    "        \"taskname\": \"intent_and_slot\",\n",
    "        \"prompt_template\": \"<|VIRTUAL_PROMPT_0|> Predict intent and slot <|VIRTUAL_PROMPT_1|> :\\n{utterance}{label}\",\n",
    "        \"total_virtual_tokens\": 10,\n",
    "        \"virtual_token_splits\": [7, 3],\n",
    "        \"truncate_field\": None,\n",
    "        \"answer_only_loss\": False,\n",
    "    },\n",
    "    {\n",
    "        \"taskname\": \"squad\",\n",
    "        \"prompt_template\": \"<|VIRTUAL_PROMPT_0|> Context: {context}\\n\\nQuestion: {question}\\n\\nAnswer:{answer}\",\n",
    "        \"total_virtual_tokens\": 15,\n",
    "        \"virtual_token_splits\": [15],\n",
    "        \"truncate_field\": \"context\",\n",
    "        \"answer_only_loss\": True,\n",
    "        \"answer_field\": \"answer\",\n",
    "    },\n",
    "    \n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dcc438b5",
   "metadata": {},
   "source": [
    "Note each `task_template` item has 5 fields. \n",
    "\n",
    "- **`prompt_template`** is a string showing the model where to place virtual tokens and how to map dataset json fields to where they belong in the model prompt. \n",
    "\n",
    "\n",
    "- **`taskname`** refers to the same `taskname` in the dataset json objects. \n",
    "\n",
    "\n",
    "- **`total_virtual_tokens`** specifies the total number of virtual tokens that will be inserted into the model prompt.\n",
    "\n",
    "\n",
    "- **`virtual_token_splits`** specifies the number of virtual tokens that belong at each `<|VIRTUAL_PROMPT_#|>` marker. `virtual_token_splits` values should add up to `total_virtual_tokens`. The number of `virtual_token_splits` should match the number of `<|VIRTUAL_PROMPT_#|>` markers. \n",
    "\n",
    "\n",
    "- **`truncate_field`** specifies which field in the data json to truncate if the length of the input exceeds the maximum sequence length of the model. If `truncate_field` is set to `None`, examples that are too long are simply dropped from the dataset.\n",
    "\n",
    "\n",
    "- **`answer_only_loss`** Whether to limit loss calculation to only the answer portion of the prompt during tuning. `True` Strongly recommended for long prompts, but shorter prompts with single word answers seem to benefit from setting this to `False`. \n",
    "\n",
    "\n",
    "- **`answer_field`** The field in the data json corresponding to the answer. The loss will only be calculated on this portion of the prompt if `answer_only_loss` is `True`. The answer field must be at the end of the prompt template.\n",
    "\n",
    "In the `task_templates` we set above, `squad` has a different number of virtual tokens than `sentiment` and `intent_and_slot`. This is because we will be p-tuning on `squad` after we p-tune on the other two tasks and **we do not need to use the same number of virtual tokens between sessions**. We also set the `truncate` field for squad because the context can sometimes be longer than the model's max sequence length, and we want that field to be truncated if the example is too long. Lastly, we set `answer_only_loss` to true for `squad` due to the longer prompt. We've found `answer_only_loss=True` to work significantly better for this task."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84579c7a",
   "metadata": {},
   "source": [
    "### Setting New Tasks\n",
    "After you p-tune your model this time, you can always go back and p-tune or prompt-tune your model on more tasks without over writting the virtual prompts who've trained this time. You can also use a different number of `total_virtual_tokens` between each training session as long as tasks ptuned or prompt tuned at the same time have the same number of `total_virtual_tokens`. For this reason, when you p-tune on a new task, you need to tell your model which of your tasks are new and which ones already exist (and thus you don't want to tune them). \n",
    "\n",
    "You do this by setting the `new_tasks` and `existing_tasks` values in the config file. Because we are p-tuning a model with no existing tasks, you should set `existing_tasks=[]` and `new_tasks=[\"sentiment\", \"intent_and_slot\"]` as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57a73e01",
   "metadata": {},
   "outputs": [],
   "source": [
    "config.model.existing_tasks = []\n",
    "config.model.new_tasks = [\"sentiment\", \"intent_and_slot\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b77e88c",
   "metadata": {},
   "source": [
    "After p-tuning and/or prompt tuning is complete, you can run inference on all tasks at the same time, regradless of their `total_virtual_tokens` value."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0d5017e",
   "metadata": {},
   "source": [
    "### Setting The Pre-Trained GPT Model\n",
    "We still need to set which GPT model we want to p-tune/prompt tune. Prompt learning methods work best with large GPT language models (5B or above), but the purposes of this tutorial, we are going to download a 345M parameter GPT model from NVIDIA NGC."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48cdf868",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check what GPT .nemo models we have available on NGC\n",
    "from nemo.collections.nlp.models.language_modeling.megatron_gpt_model import MegatronGPTModel\n",
    "MegatronGPTModel.list_available_models()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ede350ed",
   "metadata": {},
   "source": [
    "If we wanted to use the GPT model class directly, we could instantiate a trainer then download the model by calling running \n",
    "`gpt_model = MegatronGPTModel.from_pretrained(model_name=\"megatron_gpt_345m\", trainer=trainer).cuda()`. But we just need the `.nemo` file in our working NeMo directory in this tutorial, so we will download it using `wget`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "364439a1",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Download the model from NGC\n",
    "gpt_file_name = \"megatron_gpt_345m.nemo\"\n",
    "!wget --content-disposition https://api.ngc.nvidia.com/v2/models/nvidia/nemo/megatron_gpt_345m/versions/1/files/megatron_gpt_345m.nemo -O {NEMO_DIR}/{gpt_file_name}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d6a8a67",
   "metadata": {},
   "source": [
    "Now that we have a `.nemo` GPT file to work with. We need to add its path in our prompt learning config. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2778a5fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set GPT model path on prompt learning config\n",
    "config.model.language_model_path = gpt_file_name"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "943a9c83",
   "metadata": {},
   "source": [
    "We can also set where we want the final prompt tuned model to be saved by setting `model.nemo_path`. By default the tuned prompt learning model will be saved in your current working directory to a `.nemo` file with the same name as your experiment (`config.name`). Let's change the save name to be `multitask_p_tuned_gpt.nemo`. **Your model path must end in `.nemo`.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a278cbdf",
   "metadata": {},
   "outputs": [],
   "source": [
    "config.model.nemo_path = \"multitask_p_tuned_gpt.nemo\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "378a73e7",
   "metadata": {},
   "source": [
    "### Setting P-Tuning Specific Params\n",
    "Within the config file, p-tuning and prompt-tuning each have a couple of hyperparameters specific to them. We first need to tell the model that we want to do p-tuning, not prompt-tuning. To do this, we set the **`model.virtual_prompt_style`** hyperparameter like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68763763",
   "metadata": {},
   "outputs": [],
   "source": [
    "from nemo.collections.nlp.modules.common import VirtualPromptStyle\n",
    "config.model.virtual_prompt_style = VirtualPromptStyle.P_TUNING"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "947dec63",
   "metadata": {},
   "source": [
    "Then we can set the 2 p-tuning specific parameters. Reminder, pp-tuning uses an LSTM prompt encoder to predict virtual tokens. \n",
    "\n",
    "- **`p_tuning.dropout`** the LSTM prompt encoder dropout probability \n",
    "- **`p_tuning.num_layers`** the number of LSTM layers you want your p-tuning prompt encoder to have\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03f893ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "config.model.p_tuning.dropout = 0.0\n",
    "config.model.p_tuning.num_layers = 2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a988d16e",
   "metadata": {},
   "source": [
    "Let's have a look at all the values we've set in the model config. You can change any of these values in the same manner we've been using above. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12a37ada",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Final model config\n",
    "print(OmegaConf.to_yaml(config.model))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b4bc7f3",
   "metadata": {},
   "source": [
    "### Setting Prompt-Tuning Specific Params\n",
    "\n",
    "Though we are not using prompt tuning in this training session, lets go over the prompt tuning specific parameters we would use if we were. \n",
    "\n",
    "- **`prompt_tuning.new_prompt_init_methods`** Whether you want to initialize virtual token embeddings from the embeddings of existing parts of the model's vocabulary (either 'text' or 'random')\n",
    "- **`prompt_tuning.new_prompt_init_text`** The text you want to use if you have 'text' in the list above, should be None otherwise. \n",
    "\n",
    "Each of the above hyperparameters are a list of strings. \n",
    "\n",
    "`new_prompt_init_methods` would look like `[\"text\", \"random\", \"text\", \"text\"]` if you were prompt tuning on 4 tasks at once, and you wanted the second task in `new_tasks` to use random initialization. \n",
    "\n",
    "`new_prompt_init_text` might look like `[\"some text I want to use\", None, \"some other text\", \"task text goes here\"]` for those four new tasks. \n",
    "\n",
    "The order of both should correspond to the order of the tasks you have listed in `model.new_tasks`. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c048852",
   "metadata": {},
   "source": [
    "# Building the PyTorch Lightning Trainer\n",
    "NeMo models are primarily PyTorch Lightning modules - and therefore are entirely compatible with the PyTorch Lightning ecosystem.\n",
    "\n",
    "Let's first instantiate a Trainer object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90f85b2a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import pytorch_lightning as pl\n",
    "from nemo.collections.nlp.parts.nlp_overrides import NLPDDPPlugin\n",
    "from pytorch_lightning.plugins.environments.torchelastic_environment import TorchElasticEnvironment\n",
    "\n",
    "# lets modify some trainer configs\n",
    "# checks if we have GPU available and uses it\n",
    "accelerator = 'gpu' if torch.cuda.is_available() else 'cpu'\n",
    "config.trainer.accelerator = accelerator\n",
    "config.trainer.devices = 1\n",
    "config.trainer.max_epochs = 10\n",
    "config.trainer.val_check_interval = 1.0\n",
    "\n",
    "# for PyTorch Native AMP set precision=16\n",
    "config.trainer.precision = 16 if torch.cuda.is_available() else 32\n",
    "\n",
    "# remove distributed training flags\n",
    "config.trainer.strategy = None\n",
    "\n",
    "# setup cluster environment parameters\"\n",
    "# use torch elastic cluster environment so `create_process_externally` is True\n",
    "# the launcher is set to None. It will not try to spawn new processes.\n",
    "# It won't create the misconfiguration error because of the `interactive session`\n",
    "os.environ[\"LOCAL_RANK\"] = '0'\n",
    "os.environ[\"RANK\"] = '0'\n",
    "os.environ[\"WORLD_SIZE\"] = '1'\n",
    "\n",
    "plugins = [NLPDDPPlugin(find_unused_parameters=False), TorchElasticEnvironment()]\n",
    "trainer = pl.Trainer(plugins=plugins, **config.trainer)\n",
    "\n",
    "print(\"Trainer config - \\n\")\n",
    "print(OmegaConf.to_yaml(config.trainer))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d0124c1",
   "metadata": {},
   "source": [
    "# Setting up a NeMo Experiment\n",
    "\n",
    "NeMo has an experiment manager that handles logging and checkpointing for us, so let's use it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2c943ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "from nemo.utils.exp_manager import exp_manager\n",
    "\n",
    "# Set name of the experiment \n",
    "config.name = 'sentiment_intent_slot_p_tuning'\n",
    "config.exp_manager.resume_if_exists = False\n",
    "\n",
    "# Init the experiment manager and view the exp_dir\n",
    "exp_dir = exp_manager(trainer, config.get(\"exp_manager\", None))\n",
    "exp_dir = str(exp_dir)\n",
    "print(exp_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5860bd90",
   "metadata": {},
   "source": [
    "We can also set learning hyperparameters as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c4ec542",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set some of the learning parameters\n",
    "config.model.optim.lr = 1e-4\n",
    "config.model.batch_size = 16"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "298b3dce",
   "metadata": {},
   "source": [
    "# First P-Tuning Session\n",
    "The only thing left to do is load up the model and begin p-tuning!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4bda19b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from nemo.collections.nlp.models.language_modeling.megatron_gpt_prompt_learning_model import MegatronGPTPromptLearningModel\n",
    "\n",
    "model = MegatronGPTPromptLearningModel(cfg=config.model, trainer=trainer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d99f433",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Training set to 10 epochs by default in a cell above\n",
    "# Each epoch will take around 1min 15sec, but training time can vary\n",
    "trainer.fit(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d5a5798",
   "metadata": {},
   "source": [
    "When training completes, p-tuned virtual tokens from the prompt encoder are automatically moved to a `prompt_table` where all prompt tuned and p-tuned soft prompts are stored. The LSTM `prompt_encoder` is then removed from the model. This allows us to preserve previously p-tuned soft prompts while still maintaining the ability to add new p-tuned or prompt-tuned soft prompts in the future. The `prompt_table` uses the `taskname` as a key to look up the correct virtual tokens for a specified task."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6aab09d4",
   "metadata": {},
   "source": [
    "# Inference After First P-Tuning Session\n",
    "One way to run inference after p-tuning or prompt-tuning your model is to call `model.generate()`. `model.generate()` takes in \n",
    "\n",
    "- `inputs` which can be either a list of dictionary objects or `.jsonl` files containing dictionary objects, \n",
    "- `length_params`\n",
    "- `sampling_params`\n",
    "\n",
    "as arguments. More information about the [text generation API can be found here](https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/nlp/modules/common/transformer/text_generation.py).\n",
    "\n",
    "If `length_params` and `sampling_params` are set to `None`, the model generates output with a greedy decoding strategy and generates up to `30` new tokens. Most predictive downstream tasks (not text generation tasks), use greedy sampling. To see other ways to run inference with your prompt learning model and more details on how to define various inference parameters, visit `examples/nlp/language_modeling/megatron_gpt_eval.py`.\n",
    "\n",
    "Below are some randomly selected test examples from the sentiment classification and intent and slot classification test files. Notice that the `label` field is dropped from all test examples. The `MegatronPromptLearningDataset` called within `.generate()` automatically leaves fields in the prompt template empty when they are not provided in the data json. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc95e764",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_examples = [\n",
    "    {\"taskname\": \"intent_and_slot\", \"utterance\": \"tell me who will win the next presidential election\"},\n",
    "    {\"taskname\": \"intent_and_slot\", \"utterance\": \"i would like to pickup a veggie sub with a cookie from subway\"},\n",
    "    {\"taskname\": \"intent_and_slot\", \"utterance\": \"email happy new year to john\"},\n",
    "    {\"taskname\": \"intent_and_slot\", \"utterance\": \"set the alarm to seven am for work\"},\n",
    "    {\"taskname\": \"sentiment\", \"sentence\": \"The products have a low salt and fat content .\"},\n",
    "    {\"taskname\": \"sentiment\", \"sentence\": \"The agreement is valid for four years .\"},\n",
    "    {\"taskname\": \"sentiment\", \"sentence\": \"Diluted EPS rose to EUR3 .68 from EUR0 .50 .\"},\n",
    "    {\"taskname\": \"sentiment\", \"sentence\": \"The company is well positioned in Brazil and Uruguay .\"},\n",
    "    {\"taskname\": \"sentiment\", \"sentence\": \"Profit before taxes decreased by 9 % to EUR 187.8 mn in the first nine months of 2008 , compared to EUR 207.1 mn a year earlier .\"},\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3eef64f",
   "metadata": {},
   "source": [
    "This allows us to prompt the p-tuned GPT model as follows:\n",
    "\n",
    "```\n",
    "\"<|VIRTUAL_PROMPT_0|> Predict intent and slot <|VIRTUAL_PROMPT_1|> :\\nplease will you check it.\"\n",
    "\"<|VIRTUAL_PROMPT_0|> Predict intent and slot <|VIRTUAL_PROMPT_1|> :\\nset the alarm to seven am for work\"\n",
    "------------------\n",
    "\"<|VIRTUAL_PROMPT_0|> The products have a low salt and fat content . sentiment:\"\n",
    "\"<|VIRTUAL_PROMPT_0|> The agreement is valid for four years . sentiment:\"\n",
    "\n",
    "```\n",
    "\n",
    "With the correct virtual tokens inserted at each `<|VIRTUAL_PROMPT_#|>`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74a5a358",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = model.generate(inputs=test_examples, length_params=None)\n",
    "\n",
    "print('The prediction results of some sample queries with the trained model:')\n",
    "for result in response['sentences']:\n",
    "    print(result)\n",
    "    print(\"-\" * 30)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea6a533a",
   "metadata": {},
   "source": [
    "# Adding a New Task to a Previously Tuned Model\n",
    "Now that we've p-tuned our GPT model on intent/slot classification and sentiment analysis, lets add SQuAD question answering using p-tuning! First we need to update the config for the new task. \n",
    "\n",
    "# Updating The Model Config\n",
    "We need to update:\n",
    "\n",
    "1. `name`\n",
    "3. `model.restore_path`\n",
    "5. `model.existing_tasks`\n",
    "6. `model.new_tasks`\n",
    "7. `model.data.train_ds`\n",
    "8. `model.data.validation_ds`\n",
    "\n",
    "Remember that we already set `task_templates` for SQuAD when we were defining the task template for the other two tasks. We would add it here if we had not already set it above."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6adb09a3",
   "metadata": {},
   "source": [
    "Here we tell the config that we want to **load the previously p-tuned model and add new tasks to it.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5ec279d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Change the experiment name\n",
    "config.name = 'squad_p_tuning'\n",
    "\n",
    "# Change restore path from null to the p-tuned model we just finished training\n",
    "config.model.restore_path = \"multitask_p_tuned_gpt.nemo\"\n",
    "\n",
    "# Move the tasks you just p-tuned your model on to existing tasks, and add squad to the new task list\n",
    "config.model.existing_tasks = [\"sentiment\", \"intent_and_slot\"]\n",
    "config.model.new_tasks = [\"squad\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c25c3f15",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Update the dataset list to squad train and val\n",
    "# Using a subset of the training data for the sake of time\n",
    "config.model.data.train_ds = [f\"{SQUAD_DIR}/squad_short_train.jsonl\"] \n",
    "config.model.data.validation_ds = [f\"{SQUAD_DIR}/squad_val.jsonl\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "edb3097d",
   "metadata": {},
   "source": [
    "# Second P-Tuning Session"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe9c21da",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reset some model trainer and training params\n",
    "config.trainer.max_epochs = 1\n",
    "config.trainer.val_check_interval = 1000\n",
    "\n",
    "# Limiting the number of validation batches for sake of time\n",
    "config.trainer.limit_val_batches = 100\n",
    "\n",
    "config.model.optim.lr = 5e-4\n",
    "config.model.optim.sched.min_lr = 1e-5\n",
    "config.model.batch_size = 4\n",
    "\n",
    "# Reset the trainer\n",
    "plugins = [NLPDDPPlugin(find_unused_parameters=False), TorchElasticEnvironment()]\n",
    "trainer = pl.Trainer(plugins=plugins, **config.trainer)\n",
    "\n",
    "print(\"Trainer config - \\n\")\n",
    "print(OmegaConf.to_yaml(config.trainer))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6a6ad05",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reset experiment manager\n",
    "exp_dir = exp_manager(trainer, config.get(\"exp_manager\", None))\n",
    "exp_dir = str(exp_dir)\n",
    "print(exp_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3bf68b2",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Restore previously tuned model with updated config\n",
    "model = MegatronGPTPromptLearningModel.restore_from(config.model.restore_path, config.model, trainer=trainer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b3d95f1",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Prompt tune your model on squad\n",
    "# This will take around 10 min per epoch, timing is variable\n",
    "trainer.fit(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a85b2333",
   "metadata": {},
   "source": [
    "# Inference After Second P-Tuning Session\n",
    "\n",
    "Now we can run inference on all 3 tasks at once. The answers for the intent/slot and sentiment tasks should be identical to the ones from before p-tuning on squad. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7afc3d09",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Test examples with squad examples added\n",
    "test_examples = [\n",
    "    {\"taskname\": \"intent_and_slot\", \"utterance\": \"tell me who will win the next presidential election\"},\n",
    "    {\"taskname\": \"intent_and_slot\", \"utterance\": \"i would like to pickup a veggie sub with a cookie from subway\"},\n",
    "    {\"taskname\": \"intent_and_slot\", \"utterance\": \"email happy new year to john\"},\n",
    "    {\"taskname\": \"intent_and_slot\", \"utterance\": \"set the alarm to seven am for work\"},\n",
    "    {\"taskname\": \"sentiment\", \"sentence\": \"The products have a low salt and fat content .\"},\n",
    "    {\"taskname\": \"sentiment\", \"sentence\": \"The agreement is valid for four years .\"},\n",
    "    {\"taskname\": \"sentiment\", \"sentence\": \"Diluted EPS rose to EUR3 .68 from EUR0 .50 .\"},\n",
    "    {\"taskname\": \"sentiment\", \"sentence\": \"The company is well positioned in Brazil and Uruguay .\"},\n",
    "    {\"taskname\": \"sentiment\", \"sentence\": \"Profit before taxes decreased by 9 % to EUR 187.8 mn in the first nine months of 2008 , compared to EUR 207.1 mn a year earlier .\"},\n",
    "    {\"taskname\": \"squad\", \"context\": \"The build was released for download later in the day in standard 32-bit and 64-bit versions, plus a special 64-bit version which included SDKs and developer tools (Visual Studio Express and Expression Blend) for developing Metro-style apps. The Windows Store was announced during the presentation, but was not available in this build. According to Microsoft, there were about 535,000 downloads of the developer preview within the first 12 hours of its release. Originally set to expire on March 11, 2012, in February 2012 the Developer Preview's expiry date was changed to January 15, 2013.\", \"question\": \"When was the Developer preview initially intended to expire?\"},\n",
    "    {\"taskname\": \"squad\", \"context\": \"The structures of most federal governments incorporate mechanisms to protect the rights of component states. One method, known as 'intrastate federalism', is to directly represent the governments of component states in federal political institutions. Where a federation has a bicameral legislature the upper house is often used to represent the component states while the lower house represents the people of the nation as a whole. A federal upper house may be based on a special scheme of apportionment, as is the case in the senates of the United States and Australia, where each state is represented by an equal number of senators irrespective of the size of its population.\", \"question\": \"What is a bicameral legislature?\"},\n",
    "    {\"taskname\": \"squad\", \"context\": \"Imported mystery religions, which offered initiates salvation in the afterlife, were a matter of personal choice for an individual, practiced in addition to carrying on one's family rites and participating in public religion. The mysteries, however, involved exclusive oaths and secrecy, conditions that conservative Romans viewed with suspicion as characteristic of \\\"magic\\\", conspiratorial (coniuratio), or subversive activity. Sporadic and sometimes brutal attempts were made to suppress religionists who seemed to threaten traditional morality and unity, as with the senate's efforts to restrict the Bacchanals in 186 BC.\", \"question\": \"What was the practice of religion to the Romans?\"}\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "580d42cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = model.generate(inputs=test_examples, length_params=None)\n",
    "\n",
    "print('The prediction results of some sample queries with the trained model:')\n",
    "for result in response['sentences']:\n",
    "    print(result)\n",
    "    print(\"-\" * 30)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3bac402",
   "metadata": {},
   "source": [
    "For squad, remember we only trained our model on ~29% of the training examples (20k instead of ~70k) and for only 1 epoch. Results will improve if the full training set is used and the model is tuned for more training steps.\n",
    "\n",
    "This concludes our tutorial! For command line and script usage demos, [please see our docs](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/nlp/prompt_learning.html) "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
